import matplotlib.pyplot as plt
import re
import numpy as np
from scipy.interpolate import make_interp_spline
import os
import time

timestamp = time.strftime("%Y%m%d_%H%M%S")
# 读取日志文件
log_file_path = 'activation_functions.log'

# 提取训练损失数据
with open(log_file_path, 'r') as log_file:
    log_data = log_file.readlines()[30:60]  #可修改log文件数据范围

# 存储每种激活函数的训练损失数据
data = {
    "ReLU": [],
    "Sigmoid": [],
    "Tanh": []
}
accur_data = {
    "ReLU": [],
    "Sigmoid": [],
    "Tanh": []
}

# 使用正则表达式解析日志
pattern = re.compile(r'INFO:root:(\w+), epoch: (\d+), step: \s*(\d+), train_loss: ([\d.]+), test_accuracy: ([\d.]+)')

for line in log_data:
    match = pattern.search(line)
    if match:
        activation = match.group(1)  # 提取激活函数
        epoch = int(match.group(2))  # 提取epoch
        train_loss = float(match.group(4))  # 提取train_loss
        test_accuracy = float(match.group(5))  # 提取test_accuracy
        data[activation].append((epoch, train_loss))
        accur_data[activation].append((epoch, test_accuracy)) #准确率

# 打印提取的数据
#print("Data collected:")
#for activation in data:
 #   print(f"{activation}: {data[activation]}")

# 准备绘图数据------train_loss收敛画图
plot_data = []
for activation in data:
    if data[activation]:  # 确保不为空
        epochs, losses = zip(*data[activation])
        
        # 计算新横坐标值
        new_x = [1 + (i * 0.5) for i in range(len(epochs))]
        
        # 使用插值方法平滑曲线
        x_new = np.linspace(min(new_x), max(new_x), 300)  # 生成更多的x值
        spline = make_interp_spline(new_x, losses, k=3)  # k为插值多项式的阶数，3为三次样条
        losses_smooth = spline(x_new)
        
        plot_data.append((activation, x_new, losses_smooth))
    #else:
        #print(f"No data available for {activation}")  # 调试输出

# 绘图
plt.figure(figsize=(10, 6))

# 绘制每个激活函数的训练损失
for activation, x_new, losses_smooth in plot_data:
    plt.plot(x_new, losses_smooth, label=activation)

# 添加标签和标题
plt.xlabel("Epoch")
plt.ylabel("Train Loss")
plt.title("Train Loss vs Epoch for Different Activation Functions")
plt.xticks([1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5, 5.5])  # 根据需要设置x轴刻度
plt.legend()
plt.grid()
loss_folder = 'loss_plots'
os.makedirs(loss_folder,exist_ok=True)
plt.savefig(os.path.join(loss_folder, f'train_loss_plot_{timestamp}.png'))
plt.close()


# 准备绘图数据------test_accuracy收敛画图
plot_data_acc = []
for activation in accur_data:
    if accur_data[activation]:  # 确保不为空
        epochs, accuracies = zip(*accur_data[activation])
        
        # 计算新横坐标值
        new_x = [1 + (i * 0.5) for i in range(len(epochs))]
        
        # 使用插值方法平滑曲线
        x_new = np.linspace(min(new_x), max(new_x), 300)  # 生成更多的x值
        spline = make_interp_spline(new_x, accuracies, k=3)  # k为插值多项式的阶数，3为三次样条
        accuracies_smooth = spline(x_new)

        plot_data_acc.append((activation, x_new, accuracies_smooth))

# 绘制测试准确率图
plt.figure(figsize=(10, 6))
for activation, x_new, accuracies_smooth in plot_data_acc:
    plt.plot(x_new, accuracies_smooth, label=activation)

# 添加标签和标题
plt.xlabel("Epoch")
plt.ylabel("Test Accuracy")
plt.title("Test Accuracy vs Epoch for Different Activation Functions")
plt.xticks([1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5, 5.5])  # 设置x轴刻度
plt.legend()
plt.grid()

# 保存测试准确率图
accuracy_folder = 'accuracy_plots'
os.makedirs(accuracy_folder, exist_ok=True)  # 如果文件夹不存在，创建它
plt.savefig(os.path.join(accuracy_folder, f'test_accuracy_plot_{timestamp}.png'))
plt.close()  # 关闭当前图形
